import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris

# 加载示例数据 (Iris 数据集)
data = load_iris()
X = data.data
y = data.target

# 获取类的数量及其特征数量
class_labels = np.unique(y)
n_classes = len(class_labels)
n_features = X.shape[1]

# 计算整体均值
mean_overall = np.mean(X, axis=0)

# 初始化类内散度矩阵 Sw 和类间散度矩阵 Sb
Sw = np.zeros((n_features, n_features))
Sb = np.zeros((n_features, n_features))

# 计算每个类的均值和类内/类间散度矩阵
for label in class_labels:
    X_class = X[y == label]
    mean_class = np.mean(X_class, axis=0)

    # 类内散度
    Sw += np.cov(X_class.T) * (X_class.shape[0] - 1)

    # 类间散度
    n_class = X_class.shape[0]
    mean_diff = (mean_class - mean_overall).reshape(n_features, 1)
    Sb += n_class * (mean_diff).dot(mean_diff.T)

# 计算特征值和特征向量
# 计算Sw的逆矩阵与Sb的点积
# np.linalg.inv(Sw) 计算Sw的逆矩阵
# .dot(Sb) 计算逆矩阵与Sb的点积
# np.linalg.eig() 计算特征值和特征向量
eigenvalues, eigenvectors = np.linalg.eig(np.linalg.inv(Sw).dot(Sb))

# 将特征值和特征向量按特征值从大到小排序
sorted_indices = np.argsort(eigenvalues)[::-1]
sorted_eigenvalues = eigenvalues[sorted_indices]
sorted_eigenvectors = eigenvectors[:, sorted_indices]

# 选择前 k 个特征向量 (k = n_classes - 1)
k = n_classes - 1
W = sorted_eigenvectors[:, :k]

# 投影数据到新的特征空间
X_lda = X.dot(W)

# 打印结果
print('降维后的数据:\n', X_lda)

# 可视化降维后的数据 (如果选择了 2 个特征)
if k == 2:
    plt.scatter(X_lda[:, 0], X_lda[:, 1], c=y, edgecolor='k', cmap='viridis')
    plt.title('LDA Result')
    plt.xlabel('LD 1')
    plt.ylabel('LD 2')
    plt.grid()
    plt.show()